﻿' OcrProcessor.vb
Imports System.Collections.Concurrent
Imports System.Drawing
Imports System.IO
Imports System.Net
Imports System.Text
Imports System.Windows.Forms.VisualStyles.VisualStyleElement
Imports DocumentFormat.OpenXml.Vml
Imports Tesseract

Public Class OcrProcessor
    ' 新增：tessdata 路径字段
    Private ReadOnly _tessDataPath As String = System.IO.Path.Combine(AppDomain.CurrentDomain.BaseDirectory, "tessdata")

    Private ReadOnly _tesseractEngine As TesseractEngine
    Private ReadOnly _supportedImageExtensions As String() = {".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".gif"}
    'Private ReadOnly ImageExtensions As HashSet(Of String) =
    'New HashSet(Of String)(StringComparer.OrdinalIgnoreCase) From {
    '    ".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".gif"
    '}

    ' Cache the engine for reuse
    Private Shared _sharedEngine As TesseractEngine

    Private _cachedResults As New Dictionary(Of String, Tuple(Of String, DateTime))
    Private _cacheTimeout As TimeSpan = TimeSpan.FromHours(1)

    Public Sub New()
        Try
            ' 确保训练数据文件夹存在
            Dim tessDataPath = System.IO.Path.Combine(AppDomain.CurrentDomain.BaseDirectory, "tessdata")
            If Not Directory.Exists(tessDataPath) Then
                Directory.CreateDirectory(tessDataPath)
                ' 实际应用中需要确保tessdata目录中有语言数据文件
                ' 用户需要从 https://github.com/tesseract-ocr/tessdata 下载语言包
                ' 至少需要下载eng.traineddata (英文)和chi_sim.traineddata (简体中文)
            End If

            '' 初始化Tesseract引擎
            '_tesseractEngine = New TesseractEngine(tessDataPath, "eng+chi_sim", EngineMode.Default)
            '' 设置OCR引擎参数以提高准确度
            '' _tesseractEngine.SetVariable("tessedit_char_whitelist", "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ,.;:!?-_+=()[]{}@#$%^&*/'\" \|<> "")
            ''_tesseractEngine.SetVariable("tessedit_char_whitelist", "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ,.;:!?-_+=()[]{}@#$%^&*/'""\ |<> ")

            '' 对于中文识别，不应设置char_whitelist，而是根据需要调整其他参数
            '_tesseractEngine.SetVariable("preserve_interword_spaces", "0")
            '_tesseractEngine.SetVariable("language_model_penalty_non_dict_word", "0.5")
            '_tesseractEngine.SetVariable("language_model_penalty_non_freq_dict_word", "0.1")
            'If _sharedEngine Is Nothing Then
            '    ' Initialize once
            '    _sharedEngine = New TesseractEngine(tessDataPath, "eng+chi_sim", EngineMode.Default)
            '    ' Configure for speed
            '    _sharedEngine.SetVariable("tessedit_pageseg_mode", "1") ' Automatic page segmentation with OSD
            '    '_sharedEngine.SetVariable("load_system_dawg", "0") ' Disable dictionary lookups
            '    '_sharedEngine.SetVariable("load_freq_dawg", "0")
            '    '_sharedEngine.SetVariable("debug_file", "/dev/null")
            '    _sharedEngine.SetVariable("thresholding_method", "0") ' Use faster thresholding
            '    _sharedEngine.SetVariable("tessedit_create_hocr", "0") ' Disable HOCR output
            '    _sharedEngine.SetVariable("tessedit_create_txt", "1") ' Enable text output only
            'End If
            'If Environment.OSVersion.Platform = PlatformID.Win32NT Then
            '    _sharedEngine.SetVariable(“debug_file”, “NUL”)
            'Else
            '    _sharedEngine.SetVariable(“debug_file”, “/dev/null”)
            'End If
            '_sharedEngine.DefaultPageSegMode = PageSegMode.Auto
            '_tesseractEngine = _sharedEngine
        Catch ex As Exception
            Throw New Exception("初始化OCR引擎失败: " & ex.Message, ex)
        End Try
    End Sub

    Private Function CreateEngine(langs As String, Optional psm As PageSegMode = PageSegMode.Auto) As TesseractEngine
        Dim engine = New TesseractEngine(_tessDataPath, langs, EngineMode.Default)
        engine.DefaultPageSegMode = psm
        engine.SetVariable("load_system_dawg", "0")
        engine.SetVariable("load_freq_dawg", "0")
        ' Windows 下不要用 /dev/null
        If Environment.OSVersion.Platform = PlatformID.Win32NT Then
            engine.SetVariable("debug_file", "NUL")
        Else
            engine.SetVariable("debug_file", "/dev/null")
        End If
        Return engine
    End Function

    'Public Function ExtractTextFromMultipleImages(imagePaths As String()) As Dictionary(Of String, String)
    '    Dim results As New ConcurrentDictionary(Of String, String)()

    '    Parallel.ForEach(imagePaths,
    '                New ParallelOptions With {.MaxDegreeOfParallelism = Environment.ProcessorCount},
    '                Sub(path)
    '                    Dim text = ExtractTextFromImage(path)
    '                    results.TryAdd(path, text)
    '                End Sub)

    '    Return New Dictionary(Of String, String)(results)
    'End Function


    ' 判断文件是否为支持的图片格式
    Public Function IsSupportedImageFile(filePath As String) As Boolean
        Dim extension = System.IO.Path.GetExtension(filePath).ToLower()
        Return _supportedImageExtensions.Contains(extension)
        'Return ImageExtensions.Contains(extension)
    End Function

    Public Function ExtractTextFromImageWithTimeout(imagePath As String, timeoutMs As Integer) As String
        Dim ocrTask = Task.Run(Function() ExtractTextFromImage(imagePath))
        If ocrTask.Wait(timeoutMs) Then
            Return ocrTask.Result
        Else
            Logger.Log($"OCR 超时: {imagePath}")
            Return String.Empty
        End If
    End Function

    ' 从图片中提取文本
    Public Function ExtractTextFromImage(imagePath As String) As String
        If Not File.Exists(imagePath) OrElse Not IsSupportedImageFile(imagePath) Then
            Return String.Empty
        End If

        Try
            Using imageFile = Pix.LoadFromFile(imagePath)
                ' 检查图像是否成功加载
                If imageFile Is Nothing Then
                    Return String.Empty
                End If

                Dim processedImage = PreprocessImageForChineseOcr(imageFile)
                Dim needDisposeProcessed = Not Object.ReferenceEquals(processedImage, imageFile)

                Try
                    Dim text As String = String.Empty
                    Dim confidence As Single = 0

                    'SyncLock _engineLock
                    '    Using page = _tesseractEngine.Process(processedImage)
                    '        ' 这里是最容易崩的地方，所以必须在锁里
                    '        Dim raw = page.GetText()
                    '        text = If(raw, String.Empty).Trim()
                    '        confidence = page.GetMeanConfidence()
                    '    End Using
                    'End SyncLock

                    Using engine = CreateEngine("eng+chi_sim", PageSegMode.Auto)
                        Using page = engine.Process(processedImage)
                            Dim raw = page.GetText()
                            text = If(raw, String.Empty).Trim()
                            confidence = page.GetMeanConfidence()
                        End Using
                    End Using

                    If confidence < 0.5F Then
                        Logger.Log($"警告: OCR 置信度过低 ({confidence:P2}) - {imagePath}")
                    End If

                    ' 缓存、清理、后处理
                    _cachedResults(imagePath) = New Tuple(Of String, DateTime)(text, DateTime.Now)
                    If _cachedResults.Count > 100 Then CleanupCache()

                    text = PostprocessChineseText(text)
                    Return text
                Finally
                    If needDisposeProcessed Then
                        processedImage.Dispose()
                    End If
                End Try

                '' Check if we have a recent cached result
                'Dim fileInfo = New FileInfo(imagePath)
                'If _cachedResults.ContainsKey(imagePath) Then
                '    Dim cachedData = _cachedResults(imagePath)
                '    If cachedData.Item2 >= fileInfo.LastWriteTime Then
                '        Return cachedData.Item1
                '    End If
                'End If

                '' 图像预处理 - 对中文识别特别重要
                'Dim processedImage = PreprocessImageForChineseOcr(imageFile)

                ''' 对于大图片，可以选择调整大小以提高OCR速度
                ''Dim processedImage = imageFile
                ''If imageFile.Width > 3000 OrElse imageFile.Height > 3000 Then
                ''    Dim scale = 3000 / Math.Max(imageFile.Width, imageFile.Height)
                ''    processedImage = imageFile.Scale(scale, scale)
                ''End If

                '' 执行OCR
                'Using page = _tesseractEngine.Process(processedImage)
                '    Dim text = page.GetText().Trim()

                '    ' 如果OCR结果置信度过低，可能是低质量图片或非文本图片
                '    If page.GetMeanConfidence() < 0.5 Then
                '        Console.WriteLine($"警告: OCR 置信度过低 ({page.GetMeanConfidence():P2}) - {imagePath}")
                '    End If

                '    ' Cache the result
                '    _cachedResults(imagePath) = New Tuple(Of String, DateTime)(text, DateTime.Now)

                '    ' Cleanup old cache entries periodically
                '    If _cachedResults.Count > 100 Then
                '        CleanupCache()
                '    End If

                '    ' 后处理结果 - 处理可能的中文识别问题
                '    text = PostprocessChineseText(text)

                '    Return text
                'End Using
            End Using
        Catch ex As AccessViolationException ' 最常见的原生崩溃，记录后返回空字符串 
            Logger.Log($“OCR原生崩溃: {imagePath}, 错误: {ex.Message}”)
            Return String.Empty
        Catch ex As Exception
            Logger.Log($"OCR处理错误: {imagePath}, 错误: {ex.Message}")
            Return String.Empty
        End Try
    End Function

    Private Sub CleanupCache()
        Dim cutoffTime = DateTime.Now - _cacheTimeout
        Dim keysToRemove = _cachedResults.Where(Function(kv) kv.Value.Item2 < cutoffTime).Select(Function(kv) kv.Key).ToList()

        For Each key In keysToRemove
            _cachedResults.Remove(key)
        Next
    End Sub

    ' 针对中文文本的图像预处理
    Private Function PreprocessImageForChineseOcr(image As Pix) As Pix
        Try
            ' 1. 调整图像大小 - 中文识别通常需要更高的分辨率
            Dim processedImage = image

            '' 如果图像太小，考虑放大
            'If image.Width < 1000 OrElse image.Height < 1000 Then
            '    ' 放大小图像以提高中文识别率
            '    Dim scale = 1.5F
            '    processedImage = image.Scale(scale, scale)
            'ElseIf image.Width > 3000 OrElse image.Height > 3000 Then
            '    ' 缩小大图像以提高性能
            '    Dim scale = 3000 / Math.Max(image.Width, image.Height)
            '    processedImage = image.Scale(scale, scale)
            'End If

            ' Only scale if really necessary
            If image.Width < 800 OrElse image.Height < 800 Then
                ' Use smaller scale factor, e.g., 1.2 instead of 1.5
                Dim scale = 1.2F
                processedImage = image.Scale(scale, scale)
            ElseIf image.Width > 2000 OrElse image.Height > 2000 Then
                ' Increase threshold to 2000 instead of 3000
                Dim scale = 2000 / Math.Max(image.Width, image.Height)
                processedImage = image.Scale(scale, scale)
            Else
                ' Use original image when possible
                processedImage = image
            End If

            ' 2. 对图像应用二值化处理 - 对于中文文本识别很重要
            ' 选择适当的方法取决于图像类型
            ' 对于扫描文档或黑底白字，使用自适应二值化处理
            'processedImage = processedImage.BinarizeOtsu()

            ' 3. 降噪处理
            'processedImage = processedImage.Despeckle(Tesseract.Speckle.Normal)

            ' 4. 可选：增强对比度
            'processedImage = processedImage.AdjustContrast(1.5F)

            Return processedImage
        Catch ex As Exception
            Logger.Log($"图像预处理错误: {ex.Message}")
            ' 如果预处理失败，返回原始图像
            Return image
        End Try
    End Function

    ' 中文文本后处理
    Private Function PostprocessChineseText(text As String) As String
        If String.IsNullOrEmpty(text) Then
            Return text
        End If

        ' 1. 移除多余的空格 - 中文通常不需要单词间的空格
        text = System.Text.RegularExpressions.Regex.Replace(text, "(?<=[\u4e00-\u9fa5]) +(?=[\u4e00-\u9fa5])", "")

        ' 2. 修复常见的识别错误
        ' 示例：将经常错误识别的字符组合替换为正确形式
        ' text = text.Replace("某错误组合", "正确形式")

        ' 3. 移除孤立的标点符号
        text = System.Text.RegularExpressions.Regex.Replace(text, "^[,.;:!?，。；：！？、""'']+", "")
        text = System.Text.RegularExpressions.Regex.Replace(text, "[,.;:!?，。；：！？、""'']+$", "")

        ' 4. 合并多行 - 可选，取决于需求
        text = System.Text.RegularExpressions.Regex.Replace(text, "(?<=[^\n。！？.!?])\r?\n", "")

        Return text
    End Function

    ' 添加专门针对中文的OCR处理方法
    Public Function ExtractChineseTextFromImage(imagePath As String) As String
        If Not File.Exists(imagePath) OrElse Not IsSupportedImageFile(imagePath) Then
            Return String.Empty
        End If

        'If Not IsSupportedImageFile(imagePath) Then
        '    Return String.Empty
        'End If

        'Try
        '    ' 尝试使用仅中文模式
        '    _tesseractEngine.SetVariable("tessedit_pageseg_mode", "6")  ' 假设是一个统一的文本块
        '    _tesseractEngine.SetVariable("textord_min_linesize", "2.5") ' 更适合中文字体高度

        '    Using imageFile = Pix.LoadFromFile(imagePath)
        '        If imageFile Is Nothing Then
        '            Return String.Empty
        '        End If

        '        ' 特殊处理中文图像
        '        Dim processedImage = PreprocessImageForChineseOcr(imageFile)

        '        ' 先尝试仅使用中文识别
        '        Using page = _tesseractEngine.Process(processedImage, "chi_sim")
        '            Dim text = page.GetText().Trim()
        '            Dim confidence = page.GetMeanConfidence()

        '            ' 如果中文识别效果不好，尝试中英文混合识别
        '            If confidence < 0.4 Then
        '                Using pageMultiLang = _tesseractEngine.Process(processedImage, "chi_sim+eng")
        '                    Dim textMultiLang = pageMultiLang.GetText().Trim()
        '                    Dim confidenceMultiLang = pageMultiLang.GetMeanConfidence()

        '                    ' 使用置信度更高的结果
        '                    If confidenceMultiLang > confidence Then
        '                        text = textMultiLang
        '                        confidence = confidenceMultiLang
        '                    End If
        '                End Using
        '            End If

        '            Console.WriteLine($"中文OCR置信度: {confidence:P2}")

        '            ' 后处理结果
        '            text = PostprocessChineseText(text)

        '            Return text
        '        End Using
        '    End Using
        'Catch ex As Exception
        '    Console.WriteLine($"中文OCR处理错误: {imagePath}, 错误: {ex.Message}")
        '    Return String.Empty
        'End Try

        Try
            Using imageFile = Pix.LoadFromFile(imagePath)
                If imageFile Is Nothing Then Return String.Empty

                Dim processedImage = PreprocessImageForChineseOcr(imageFile)
                Dim needDisposeProcessed = Not Object.ReferenceEquals(processedImage, imageFile)

                Try
                    Dim bestText As String = String.Empty
                    Dim bestConf As Single = 0

                    ' 先仅中文
                    Using engineCn = CreateEngine("chi_sim", PageSegMode.SingleBlock)
                        Using page = engineCn.Process(processedImage)
                            bestText = If(page.GetText(), String.Empty).Trim()
                            bestConf = page.GetMeanConfidence()
                        End Using
                    End Using

                    ' 再尝试中英混合，取更高置信度
                    Using engineMix = CreateEngine("chi_sim+eng", PageSegMode.SingleBlock)
                        Using page2 = engineMix.Process(processedImage)
                            Dim txt2 = If(page2.GetText(), String.Empty).Trim()
                            Dim conf2 = page2.GetMeanConfidence()
                            If conf2 > bestConf Then
                                bestText = txt2
                                bestConf = conf2
                            End If
                        End Using
                    End Using

                    Logger.Log($"中文OCR置信度: {bestConf:P2}")
                    bestText = PostprocessChineseText(bestText)
                    Return bestText
                Finally
                    If needDisposeProcessed Then processedImage.Dispose()
                End Try
            End Using
        Catch ex As AccessViolationException
            Logger.Log($"中文OCR原生崩溃: {imagePath}, 错误: {ex.Message}")
            Return String.Empty
        Catch ex As Exception
            Logger.Log($"中文OCR处理错误: {imagePath}, 错误: {ex.Message}")
            Return String.Empty
        End Try
    End Function

    ' 检查OCR语言文件
    Public Function CheckOcrLanguageFiles() As Dictionary(Of String, Boolean)
        Dim results As New Dictionary(Of String, Boolean)()

        ' 检查语言文件路径
        Dim tessDataPath = System.IO.Path.Combine(AppDomain.CurrentDomain.BaseDirectory, "tessdata")
        results.Add("tessdata目录存在", Directory.Exists(tessDataPath))

        ' 检查英文语言文件
        Dim engFile = System.IO.Path.Combine(tessDataPath, "eng.traineddata")
        results.Add("英文语言文件(eng.traineddata)", File.Exists(engFile))

        ' 检查简体中文语言文件
        Dim chiSimFile = System.IO.Path.Combine(tessDataPath, "chi_sim.traineddata")
        results.Add("简体中文语言文件(chi_sim.traineddata)", File.Exists(chiSimFile))

        ' 检查繁体中文语言文件(可选)
        Dim chiTraFile = System.IO.Path.Combine(tessDataPath, "chi_tra.traineddata")
        results.Add("繁体中文语言文件(chi_tra.traineddata)", File.Exists(chiTraFile))

        ' 尝试初始化引擎
        Try
            Using engine = New TesseractEngine(tessDataPath, "chi_sim", EngineMode.Default)
                results.Add("中文引擎初始化", True)
            End Using
        Catch ex As Exception
            results.Add("中文引擎初始化", False)
            results.Add("初始化错误", ex.Message)
        End Try

        Return results
    End Function

    ' 显示OCR诊断信息
    Public Sub ShowOcrDiagnostics()
        Dim results = CheckOcrLanguageFiles()
        Dim message As New StringBuilder("OCR诊断信息:" & Environment.NewLine & Environment.NewLine)

        For Each result In results
            message.AppendLine($"{result.Key}: {(If(result.Value, "✓", "✗"))}")
        Next

        MessageBox.Show(message.ToString(), "OCR诊断", MessageBoxButtons.OK,
                  If(results.Values.All(Function(v) v), MessageBoxIcon.Information, MessageBoxIcon.Warning))
    End Sub

    ' 清理资源
    Public Sub Dispose()
        _tesseractEngine?.Dispose()
    End Sub
End Class



